# 视听事件检测模型训练代码（带详细注释）
import os
import time
import random
import json
from tqdm import tqdm  # 进度条显示库

import torch
import torch.nn as nn
from torch.nn.utils import clip_grad_norm_  # 梯度裁剪工具
from torch.utils.data import DataLoader  # 数据加载器
from tensorboardX import SummaryWriter  # TensorBoard日志记录
from torch.optim.lr_scheduler import StepLR, MultiStepLR  # 学习率调度器

import numpy as np
from configs.opts import parser  # 命令行参数解析器
from model.main_model import supv_main_model as main_model  # 主模型导入
from utils import AverageMeter, Prepare_logger, get_and_save_args  # 工具函数：平均值计算、日志配置、参数处理
from utils.Recorder import Recorder  # 实验记录器
from dataset.AVE_dataset import AVEDataset  # AVE数据集类
import torch.nn.functional as F  # torch功能函数库

# ================================= 随机种子配置 ============================
# 设置随机种子以保证实验可重复性
SEED = 43
random.seed(SEED)  # Python随机库种子
np.random.seed(seed=SEED)  # Numpy随机库种子
torch.manual_seed(seed=SEED)  # PyTorch CPU随机种子
torch.cuda.manual_seed(seed=SEED)  # PyTorch GPU随机种子
torch.backends.cudnn.deterministic = True  # 确保CUDA卷积操作确定性
torch.backends.cudnn.benchmark = False  # 禁用CUDA自动优化（保证可重复性）

# 加载模型配置文件
config_path = 'configs/main.json'
with open(config_path) as fp:
    config = json.load(fp)  # 读取JSON格式的配置参数
print(config)  # 打印配置信息
# =============================================================================

def AVPSLoss(av_simm, soft_label):
    """
    全监督设置下的视听对相似度损失函数
    对应论文中的公式(8, 9)
    参数:
        av_simm: 视听相似度分数，形状为[batch_size, 10]
        soft_label: 软标签，与av_simm形状一致
    返回:
        均方误差损失值
    """
    # 对相似度分数应用ReLU激活，过滤负值
    relu_av_simm = F.relu(av_simm)
    # 计算每个样本的相似度分数总和（保持维度）
    sum_av_simm = torch.sum(relu_av_simm, dim=-1, keepdim=True)
    # 归一化相似度分数
    avg_av_simm = relu_av_simm / (sum_av_simm + 1e-8)  # 加微小值避免除零
    # 计算均方误差损失
    loss = nn.MSELoss()(avg_av_simm, soft_label)
    return loss


def main():
    # 全局变量声明
    global args, logger, writer, dataset_configs
    # 模型性能统计变量
    global best_accuracy, best_accuracy_epoch
    best_accuracy, best_accuracy_epoch = 0, 0  # 初始化最佳准确率和对应epoch
    # 加载并保存数据集配置
    dataset_configs = get_and_save_args(parser)
    # 将数据集配置设置为解析器默认参数
    parser.set_defaults(**dataset_configs)
    # 解析命令行参数
    args = parser.parse_args()
    # 设置CUDA设备顺序和可见设备
    os.environ['CUDA_DEVICE_ORDER'] = "PCI_BUS_ID"  # 按PCI总线ID排序GPU
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu  # 指定使用的GPU编号

    '''创建快照目录，用于复制代码和保存模型'''
    if not os.path.exists(args.snapshot_pref):
        os.makedirs(args.snapshot_pref)  # 不存在则创建目录

    # 如果指定了恢复 checkpoint，使用其所在目录作为快照目录
    if os.path.isfile(args.resume):
        args.snapshot_pref = os.path.dirname(args.resume)

    # 配置日志记录器
    logger = Prepare_logger(args, eval=args.evaluate)

    # 训练模式下打印日志信息
    if not args.evaluate:
        logger.info(f'\n创建文件夹: {args.snapshot_pref}')
        logger.info('\n运行参数\n\n{}\n'.format(json.dumps(vars(args), indent=4)))
    # 评估模式下打印日志信息
    else:
        logger.info(f'\n日志文件将保存到 {args.snapshot_pref}/Eval.log.')

    '''加载数据集'''
    # 训练集数据加载器
    train_dataloader = DataLoader(
        AVEDataset('./data/', split='train'),  # 加载训练集
        batch_size=args.batch_size,  # 批次大小
        shuffle=True,  # 打乱数据
        num_workers=8,  # 多线程加载
        pin_memory=True  # 锁定内存（加速GPU数据传输）
    )

    # 测试集数据加载器
    test_dataloader = DataLoader(
        AVEDataset('./data/', split='test'),  # 加载测试集
        batch_size=args.test_batch_size,  # 测试批次大小
        shuffle=False,  # 不打乱测试数据
        num_workers=8,  # 多线程加载
        pin_memory=True  # 锁定内存
    )

    '''模型配置'''
    # 初始化主模型，传入模型配置参数
    mainModel = main_model(config['model'])
    # 多GPU并行训练
    mainModel = nn.DataParallel(mainModel).cuda()
    # 获取模型可学习参数
    learned_parameters = mainModel.parameters()
    # 初始化Adam优化器
    optimizer = torch.optim.Adam(learned_parameters, lr=args.lr)
    # 学习率调度器：多步学习率衰减，在指定epoch处调整
    scheduler = MultiStepLR(optimizer, milestones=[10, 20, 30], gamma=0.5)

    # 初始化损失函数（均移到GPU）
    criterion = nn.BCEWithLogitsLoss().cuda()  # 二分类交叉熵损失（带logits）
    criterion_event = nn.CrossEntropyLoss().cuda()  # 多分类交叉熵损失

    '''从checkpoint恢复训练'''
    if os.path.isfile(args.resume):
        logger.info(f"\n加载Checkpoint: {args.resume}\n")
        # 加载模型权重
        mainModel.load_state_dict(torch.load(args.resume))
    # 若指定了resume但文件不存在，抛出异常
    elif args.resume != "" and (not os.path.isfile(args.resume)):
        raise FileNotFoundError

    '''仅执行评估'''
    if args.evaluate:
        logger.info(f"\n开始评估..")
        # 调用验证函数进行评估
        validate_epoch(mainModel, test_dataloader, criterion, criterion_event, epoch=0, eval_only=True)
        return

    '''TensorBoard日志和代码备份'''
    # 初始化TensorBoard写入器
    writer = SummaryWriter(args.snapshot_pref)
    # 初始化实验记录器，忽略指定文件夹
    recorder = Recorder(args.snapshot_pref, ignore_folder="Exps/")
    # 记录实验参数
    recorder.writeopt(args)

    '''训练和测试主循环'''
    for epoch in range(args.n_epoch):
        # 训练一个epoch，返回平均损失
        loss = train_epoch(mainModel, train_dataloader, criterion, criterion_event, optimizer, epoch)

        # 按指定频率或最后一个epoch执行验证
        if ((epoch + 1) % args.eval_freq == 0) or (epoch == args.n_epoch - 1):
            # 验证一个epoch，返回准确率
            acc = validate_epoch(mainModel, test_dataloader, criterion, criterion_event, epoch)
            # 更新最佳准确率和对应epoch
            if acc > best_accuracy:
                best_accuracy = acc
                best_accuracy_epoch = epoch
                # 保存最佳模型
                save_checkpoint(
                    mainModel.state_dict(),
                    top1=best_accuracy,
                    task='Supervised',
                    epoch=epoch + 1,
                )
            # 打印最佳性能信息
            print("-----------------------------")
            print("最佳准确率和对应epoch:", best_accuracy, best_accuracy_epoch)
            print("-----------------------------")
        # 更新学习率
        scheduler.step()


def train_epoch(model, train_dataloader, criterion, criterion_event, optimizer, epoch):
    """
    训练一个epoch的函数
    参数:
        model: 训练的模型
        train_dataloader: 训练集数据加载器
        criterion: 二分类损失函数
        criterion_event: 事件分类损失函数
        optimizer: 优化器
        epoch: 当前epoch编号
    返回:
        整个epoch的平均损失
    """
    # 初始化性能指标记录器
    batch_time = AverageMeter()  # 批次处理时间
    data_time = AverageMeter()  # 数据加载时间
    losses = AverageMeter()  # 损失值
    train_acc = AverageMeter()  # 训练准确率
    end_time = time.time()  # 记录起始时间

    # 设置模型为训练模式
    model.train()
    # 注意：此处将模型设置为双精度（double），因为提取的特征是双精度类型
    # 这也会导致模型参数占用内存翻倍
    model.double()
    # 清空梯度
    optimizer.zero_grad()

    # 遍历训练数据集
    for n_iter, batch_data in enumerate(train_dataloader):
        # 更新数据加载时间
        data_time.update(time.time() - end_time)
        
        '''将数据输入模型'''
        # 解析批次数据：视觉特征、音频特征、标签
        visual_feature, audio_feature, labels = batch_data
        # 将标签转换为双精度并移到GPU
        labels = labels.double().cuda()
        # 模型前向传播，输出各分支结果
        is_event_scores, event_scores, audio_visual_gate, av_score = model(visual_feature, audio_feature)
        # 调整输出维度（转置并压缩冗余维度）
        is_event_scores = is_event_scores.transpose(1, 0).squeeze().contiguous()
        audio_visual_gate = audio_visual_gate.transpose(1, 0).squeeze().contiguous()

        # 处理标签：提取前景标签和事件类别标签
        labels_foreground = labels[:, :, :-1]  # [32, 10, 28]，排除背景类别
        labels_BCE, labels_evn = labels_foreground.max(-1)  # 提取二分类标签和事件类别标签
        labels_event, _ = labels_evn.max(-1)  # 提取最终事件类别标签

        # 计算各部分损失
        loss_is_event = criterion(is_event_scores, labels_BCE.cuda())  # 事件存在性损失
        label_is_gate = criterion(audio_visual_gate, labels_BCE.cuda())  # 门控损失
        loss_cas = criterion_event(av_score, labels_event.cuda())  # 视听分数分类损失
        loss_event_class = criterion_event(event_scores, labels_event.cuda())  # 事件分类损失

        # 总损失（多损失项加权求和）
        loss = loss_is_event + label_is_gate + loss_event_class + loss_cas
        # 反向传播计算梯度
        loss.backward()

        '''计算训练准确率'''
        acc = compute_accuracy_supervised(is_event_scores, event_scores, labels)
        # 更新准确率记录（按样本数加权）
        train_acc.update(acc.item(), visual_feature.size(0) * 10)

        '''梯度裁剪'''
        if args.clip_gradient is not None:
            total_norm = clip_grad_norm_(model.parameters(), args.clip_gradient)
            # 可选：打印梯度裁剪信息
            # if total_norm > args.clip_gradient:
            #     logger.info(f'Clipping gradient: {total_norm} with coef {args.clip_gradient/total_norm}.')

        '''更新模型参数'''
        optimizer.step()
        optimizer.zero_grad()  # 清空梯度

        # 更新损失记录
        losses.update(loss.item(), visual_feature.size(0) * 10)
        # 更新批次处理时间
        batch_time.update(time.time() - end_time)
        end_time = time.time()

        '''记录迭代损失到TensorBoard'''
        writer.add_scalar('Train_data/loss', losses.val, epoch * len(train_dataloader) + n_iter + 1)

        '''打印训练日志'''
        if n_iter % args.print_freq == 0:
            logger.info(
                f'Train Epoch: [{epoch}][{n_iter}/{len(train_dataloader)}]\t'
                f'Loss {losses.val:.4f} ({losses.avg:.4f})\t'
                f'Prec@1 {train_acc.val:.3f} ({train_acc.avg: .3f})'
            )

        '''记录epoch平均损失到TensorBoard'''
        writer.add_scalar('Train_epoch_data/epoch_loss', losses.avg, epoch)

    # 返回整个epoch的平均损失
    return losses.avg


@torch.no_grad()  # 禁用梯度计算（验证阶段不需要）
def validate_epoch(model, test_dataloader, criterion, criterion_event, epoch, eval_only=False):
    """
    验证/测试一个epoch的函数
    参数:
        model: 待验证的模型
        test_dataloader: 测试集数据加载器
        criterion: 二分类损失函数
        criterion_event: 事件分类损失函数
        epoch: 当前epoch编号
        eval_only: 是否仅执行评估（不记录训练相关日志）
    返回:
        整个epoch的平均准确率
    """
    # 初始化性能指标记录器
    batch_time = AverageMeter()  # 批次处理时间
    data_time = AverageMeter()  # 数据加载时间
    losses = AverageMeter()  # 损失值
    accuracy = AverageMeter()  # 准确率
    end_time = time.time()  # 记录起始时间

    # 设置模型为评估模式
    model.eval()
    # 设置模型为双精度（与训练保持一致）
    model.double()

    # 遍历测试数据集
    for n_iter, batch_data in enumerate(test_dataloader):
        # 更新数据加载时间
        data_time.update(time.time() - end_time)

        '''将数据输入模型'''
        # 解析批次数据：视觉特征、音频特征、标签
        visual_feature, audio_feature, labels = batch_data
        # 将标签转换为双精度并移到GPU
        labels = labels.double().cuda()
        # 获取批次大小
        bs = visual_feature.size(0)
        # 模型前向传播
        is_event_scores, event_scores, audio_visual_gate, _ = model(visual_feature, audio_feature)
        # 调整输出维度
        is_event_scores = is_event_scores.transpose(1, 0).squeeze()
        audio_visual_gate = audio_visual_gate.transpose(1, 0).squeeze()

        # 处理标签
        labels_foreground = labels[:, :, :-1]
        labels_BCE, labels_evn = labels_foreground.max(-1)
        labels_event, _ = labels_evn.max(-1)

        # 计算损失
        loss_is_event = criterion(is_event_scores, labels_BCE.cuda())
        loss_is_gate = criterion(audio_visual_gate, labels_BCE.cuda())
        loss_event_class = criterion_event(event_scores, labels_event.cuda())
        total_loss = loss_is_event + loss_event_class + loss_is_gate

        # 计算准确率
        acc = compute_accuracy_supervised(is_event_scores, event_scores, labels)
        # 更新准确率记录
        accuracy.update(acc.item(), bs * 10)

        # 更新批次处理时间
        batch_time.update(time.time() - end_time)
        end_time = time.time()
        # 更新损失记录
        losses.update(total_loss.item(), bs * 10)

        '''打印验证日志'''
        if n_iter % args.print_freq == 0:
            logger.info(
                f'Test Epoch [{epoch}][{n_iter}/{len(test_dataloader)}]\t'
                f'Loss {losses.val:.4f} ({losses.avg:.4f})\t'
                f'Prec@1 {accuracy.val:.3f} ({accuracy.avg:.3f})'
            )

    '''记录验证日志到TensorBoard'''
    if not eval_only:
        writer.add_scalar('Val_epoch_data/epoch_loss', losses.avg, epoch)
        writer.add_scalar('Val_epoch/Accuracy', accuracy.avg, epoch)

    # 打印最终评估结果
    logger.info(
        f'**************************************************************************\t'
        f"\t评估结果（准确率）: {accuracy.avg:.4f}%."
    )
    # 返回平均准确率
    return accuracy.avg


def compute_accuracy_supervised(is_event_scores, event_scores, labels):
    """
    计算监督训练模式下的准确率
    参数:
        is_event_scores: 事件存在性预测分数
        event_scores: 事件类别预测分数
        labels: 真实标签
    返回:
        准确率（百分比）
    """
    # 提取真实标签的类别（忽略背景类）
    _, targets = labels.max(-1)
    # 事件存在性分数转换为概率（sigmoid激活）
    is_event_scores = is_event_scores.sigmoid
    # 续：准确率计算函数完整实现
def compute_accuracy_supervised(is_event_scores, event_scores, labels):
    """
    计算监督训练模式下的准确率
    参数:
        is_event_scores: 事件存在性预测分数
        event_scores: 事件类别预测分数
        labels: 真实标签
    返回:
        准确率（百分比）
    """
    # 提取真实标签的类别（忽略背景类）
    _, targets = labels.max(-1)
    # 事件存在性分数转换为概率（sigmoid激活）
    is_event_scores = is_event_scores.sigmoid()
    # 大于0.5视为存在事件（二分类阈值）
    scores_pos_ind = is_event_scores > 0.5
    # 生成掩码：标记不存在事件的位置
    scores_mask = scores_pos_ind == 0
    # 事件类别预测：取分数最高的类别
    _, event_class = event_scores.max(-1)  # 前景事件分类结果
    # 初始化预测结果：存在事件的位置填入类别，否则为0
    pred = scores_pos_ind.long()
    pred *= event_class[:, None]  # 扩展维度并赋值类别
    # 不存在事件的位置标记为背景类（28表示背景）
    pred[scores_mask] = 28
    # 计算预测正确的样本数
    correct = pred.eq(targets)
    correct_num = correct.sum().double()
    # 计算准确率（转换为百分比）
    acc = correct_num * (100. / correct.numel())
    return acc


def save_checkpoint(state_dict, top1, task, epoch):
    """
    保存模型 checkpoint
    参数:
        state_dict: 模型状态字典
        top1: 验证集top1准确率
        task: 任务类型（如Supervised）
        epoch: 当前epoch编号
    """
    # 构建模型文件名
    model_name = f'{args.snapshot_pref}/model_epoch_{epoch}_top1_{top1:.3f}_task_{task}_best_model.pth.tar'
    # 保存模型权重
    torch.save(state_dict, model_name)


# 程序入口
if __name__ == '__main__':
    main()  # 执行主函数